//
// Created by myrica on 3/26/25.
//

#include "Pass/MoveAlloca.h"
#include <ranges>
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Support/LLVM.h"

namespace irispass {
    std::unique_ptr<mlir::Pass> createMoveAllocaToEntryPass() {
        return std::make_unique<MoveAllocaToEntryPass>();
    }

    void MoveAllocaToEntryPass::runOnOperation() {
        mlir::func::FuncOp func = getOperation();

        if (func.getBody().empty())
            return;

        mlir::Block &entryBlock = func.getBody().front();

        mlir::SmallVector<std::tuple<mlir::memref::AllocaOp, mlir::linalg::FillOp, mlir::Value> >
                alloca_fill_zero;

        llvm::errs() << "Current function: " << func.getName() << "\n";
        func.walk([&](mlir::memref::AllocaOp alloca) {
            if (alloca->getBlock() == &entryBlock)
                return;

            int fill_count = 0;
            mlir::linalg::FillOp fill = nullptr;
            mlir::Value zero = nullptr;
            for (auto user: alloca->getUsers()) {
                if (llvm::isa<mlir::linalg::FillOp>(user)) {
                    fill_count += 1;
                    fill = llvm::cast<mlir::linalg::FillOp>(user);
                }
            }
            assert(fill_count <= 1);
            if (fill) {
                assert(fill.getNumOperands() == 2);
                for (auto opr: fill->getOperands()) {
                    if (llvm::isa<mlir::arith::ConstantOp>(opr.getDefiningOp()))
                        zero = opr;
                }
            }
            alloca_fill_zero.emplace_back(alloca, fill, zero);
        });

        llvm::errs() << "Num of allocas to move: " << alloca_fill_zero.size() << "\n";
        // 逆序移动以保持原顺序
        for (auto &[alloca, fill, zero]: alloca_fill_zero) {
            auto term = entryBlock.getTerminator();

            alloca->moveBefore(term);
            if (fill) {
                zero.getDefiningOp()->moveBefore(alloca);
                fill->moveAfter(alloca);
            }
        }
    }
}
